# Copyright (c) 2008
#   Matt Pizzimenti (mjpizz@gmail.com, www.livelearncode.com)
#
# All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
# 
# Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
# 
# Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 
# Neither the name of the original author nor the names of contributors
# may be used to endorse or promote products derived from this software
# without specific prior written permission.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
import os
import os.path
import signal
import exceptions
import popen2
import paramiko
import interaction

import support
import events
from receipts import ReceiptManager as _ReceiptManager

import subprocess

class AuthenticationFailed(exceptions.Exception): pass
class UnexpectedQuitError(exceptions.Exception): pass
class FilePullFailed(exceptions.Exception): pass


class StrategyResponse(object):
    """Encapsulates the result of a command executed
    on the (possibly remote) machine.  Both stdout
    and stderr are available as newline-separated strings."""
    
    stdout = property(lambda self: self.__stdout)
    stderr = property(lambda self: self.__stderr)
    status = property(lambda self: self.__status)
    
    def __init__(self, stdout, stderr, status):
        self.__stdout = stdout
        self.__stderr = stderr
        self.__status = status


class ExecutionStrategy(object):
    """Interface for strategies used by Roles.  This abstracts
    the logic and libraries that are being used to:
    
        1. Execute shell commands
        2. Log stdin, stdout, and stderr
        3. Expect input from the user
    
    Implementations should take care to ensure that they adhere to
    the exact behavior as described in the docstrings."""
    
    def __init__(self, user):
        self.user = user
    
    def log_stdin(self, line):
        """logs a line of standard input"""
        event = events.StdinEvent(line, self.user)
        _ReceiptManager.get_instance().log_event( event )

    def log_stdout(self, line):
        """logs a line of standard output"""
        event = events.StdoutEvent(line, self.user)
        _ReceiptManager.get_instance().log_event( event )

    def log_stderr(self, line):
        """logs a line of standard error"""
        event = events.StderrEvent(line, self.user)
        _ReceiptManager.get_instance().log_event( event )

    def log_cmd(self, cmd):
        """logs a the actual shell command being executed"""
        event = events.ShellCommandEvent(cmd, self.user)
        _ReceiptManager.get_instance().log_event( event )
    
    def log_pulling_file(self, local_filename, remote_filename):
        """logs the action of pulling a local file onto the machine"""
        event = events.PullEvent(local_filename, remote_filename, self.user)
        _ReceiptManager.get_instance().log_event( event )

    def log_making_dir(self, remote_dir):
        """logs the action of making a directory on the machine"""
        event = events.MkdirEvent(remote_dir, self.user)
        _ReceiptManager.get_instance().log_event( event )

    def expect_password_for(self, pattern):
        """Prompts the user for password-style input (not echoed)
        when a prompt that looks like +pattern+ appears in the
        stdout or stderr streams."""
        raise NotImplementedError

    def expect_input_for(self, pattern):
        """Prompts the user for input when a prompt that looks
        like +pattern+ appears in the stdout or stderr streams."""
        raise NotImplementedError

    def respond_with(self, pattern, response):
        """Automaticaly responds to prompts that look like +pattern+
        with input +response+."""
        raise NotImplementedError

    def respond_with_password(self, pattern, password):
        """Automaticaly responds to prompts taht look like +pattern+ with
        a password response (should not echo password to screen)"""
        raise NotImplementedError

    def stop_expecting(self, pattern):
        """Stops expecting user input for prompts that look
        like +pattern+."""
        raise NotImplementedError
    
    def stop_responding_to(self, pattern):
        """Stops responding automatically to prompts that
        look like +pattern+."""
        raise NotImplementedError
    
    def exec_cmd(self, cmd):
        """Executes a command on the machine and returns the contents
        of the stdout stream.  Streams should be logged to the console. Input
        should be properly expect-ed and inputted to the process."""
        raise NotImplementedError
    
    def pull(self, local_filename, remote_filename):
        """Sends the given local file from the local machine to this user."""
        raise NotImplementedError


class DryRunStrategy(ExecutionStrategy):
    """strategy used to show what commands are being run, but without
    actually executing anything"""
    
    def __init__(self, user):
        super(DryRunStrategy, self).__init__(user)
    
    def expect_password_for(self, pattern): pass

    def expect_input_for(self, pattern): pass

    def respond_with(self, pattern, response): pass

    def respond_with_password(self, pattern, password): pass

    def stop_expecting(self, pattern): pass

    def stop_responding_to(self, pattern): pass

    def exec_cmd(self, cmd):
        self.log_cmd(cmd)
        return StrategyResponse(stdout = "", stderr = "", status = 0)

    def pull(self, local_filename, remote_filename):
        self.log_pulling_file(local_filename, remote_filename)
        

class ExpecterStrategy(ExecutionStrategy):
    
    def __init__(self, user):
        super(ExpecterStrategy, self).__init__(user)
        self.prompt_prefix = "[%s] " % user.fullname
        self.expecter_factory = interaction.ExpecterFactory()
    
    def expect_password_for(self, pattern):
        self.expecter_factory.add_password_responder(pattern)

    def expect_input_for(self, pattern):
        self.expecter_factory.add_responder(pattern)

    def respond_with(self, pattern, response):
        self.expecter_factory.add_auto_responder(pattern, response)

    def respond_with_password(self, pattern, password):
        self.expecter_factory.add_auto_password_responder(pattern, password)

    def stop_expecting(self, pattern):
        self.expecter_factory.remove_responder(pattern)

    def stop_responding_to(self, pattern):
        self.expecter_factory.remove_responder(pattern)

    def _watch_exec_stream_and_expect_input(self, cmd, stdin, stdout, stderr, exec_callback, kill_callback):
        
        # create an expecter based on the patterns we have populated into our Factory
        # (the expecter will log any input via the physical_machine's log_stdin)
        expecter = self.expecter_factory.create_expecter(self.prompt_prefix, stdin, self.log_stdin)

        # create stream watches that will log stdout and stderr to the console and
        # also utilize the expecter to handle input prompts
        # FIXME: Issue 18: technically, the stdout and stderr strings could be logged out-of-order (since the Monitors execute in separate threads)
        stdout_watcher = interaction.StreamMonitor(
            stream = stdout,
            stream_log_callback = self.log_stdout,
            expecter = expecter,
            description = "%s [stdout for cmd = '%s']" % (self.user.fullname, cmd),
            force_quit_callback = kill_callback
        )
        stderr_watcher = interaction.StreamMonitor(
            stream = stderr,
            stream_log_callback = self.log_stderr,
            expecter = expecter,
            description = "%s [stderr for cmd = '%s']" % (self.user.fullname, cmd),
            force_quit_callback = kill_callback
        )

        # start watching the streams
        stdout_watcher.start()
        stderr_watcher.start()

        # execute the command and get the callback
        try:
            retcode = exec_callback()
            stdout_watcher.join()
            stderr_watcher.join()
        except interaction.QuitForcedInStreamMonitor, interaction.ExceptionOccurredInStreamMonitor:
            raise UnexpectedQuitError("shell command unexpectedly killed (%s)" % sys.exc_info()[1])

        # return
        return StrategyResponse(stdout = stdout_watcher.output, stderr = stderr_watcher.output, status = retcode)


# class LocalExpecterStrategy(ExpecterStrategy):
# 
#     def __init__(self, user):
#         super(LocalExpecterStrategy, self).__init__(user)
#     
#     def exec_cmd(self, cmd):
#         
#         # TODO: use Subprocess instead?
#         # open a pipe for the command
#         self.log_cmd(cmd)
#         pipe = popen2.Popen3(cmd, True, 5) # TODO: why 5?
#     
#         # grab all the IO streams
#         stdin = pipe.tochild
#         stdout = pipe.fromchild
#         stderr = pipe.childerr
#     
#         # use the helper to monitor the command
#         def kill_this_process():
#             os.kill(pipe.pid, signal.SIGKILL)
#         return self._watch_exec_stream_and_expect_input(cmd, stdin, stdout, stderr, pipe.wait, kill_this_process)
# 	
# 	
	

class LocalExpecterStrategy(ExpecterStrategy):
	
    def __init__(self, user):
        super(LocalExpecterStrategy, self).__init__(user)
	
	
	def exec_cmd(self, cmd):
		
		self.log_cmd(cmd)
		
		#FIXME: error handling if the user does not have a shell
		shell = os.environ['SHELL']
		
		p = Popen(args = cmd,
			shell=True,
			executable=shell,
			stdin=PIPE,
			stdout=PIPE,
			stderr=PIPE,
			close_fds=False
			)
			
		stdin = p.stdin
		stdout = p.stdout
		stderr = p.stderr
		
		def kill_this_process():
			os.kill(p.pid, signal.SIGKILL)
		
		return self._watch_exec_stream_and_expect_input(cmd, stdin, stdout, stderr, p.wait, kill_this_process)
        
		
class RemoteExpecterStrategy(ExpecterStrategy):

    def __init__(self, user):
        super(RemoteExpecterStrategy, self).__init__(user)
        self.ssh_client = paramiko.SSHClient()
        self.ssh_client.load_system_host_keys()
        try:
            self.ssh_client.connect(hostname = user.host, username = user.name, password = user.password)
        except:
            raise AuthenticationFailed("couldn't log in for user '%s' (reason: %s)" % (self.user.fullname, sys.exc_info()[1]))
    
    def exec_cmd(self, cmd):
        # open an SSH channel (acts like a socket)
        channel = self.ssh_client.get_transport().open_session()
        
        # get a pseudo-terminal so that certain commands (like scp and su) behave more nicely
        channel.get_pty()

        # grab the streams on that channel
        bufsize = -1
        stdin = channel.makefile('wb', bufsize)
        stdout = channel.makefile('rb', bufsize)
        stderr = channel.makefile_stderr('rb', bufsize)

        # create a callback that will execute the command and return its status
        def exec_and_return_status_code():
            self.log_cmd(cmd)
            channel.exec_command(cmd)
            return channel.recv_exit_status()

        # use the helper to monitor the command
        def kill_this_process():
            channel.close()
        return self._watch_exec_stream_and_expect_input(cmd, stdin, stdout, stderr, exec_and_return_status_code, kill_this_process)
    
    def _mkdir(self, sftp_client, remote_dir):
        try:
            sftp_client.lstat(remote_dir)
            dir_exists = True
        except:
            dir_exists = False
            
        if dir_exists:
            raise FilePullFailed("directory '%s' already exists on the remote machine" % remote_dir)
        else:
            self.log_making_dir(remote_dir)
            try:
                sftp_client.mkdir(remote_dir)
            except IOError:
                raise FilePullFailed("failed to create remote directory '%s' (reason: %s)" % (remote_dir, sys.exc_info()[1]))
    
    def _put(self, sftp_client, local_filename, remote_filename):
        try:
            sftp_client.lstat(remote_filename)
            file_exists = True
        except:
            file_exists = False
        
        if file_exists:
            raise FilePullFailed("file '%s' already exists on the remote machine" % remote_filename)
        else:
            self.log_pulling_file(local_filename, remote_filename)
            try:
                sftp_client.put(local_filename, remote_filename)
            except IOError:
                raise FilePullFailed("failed to send file '%s' to remote location '%s' (reason: %s)" % (local_filename, remote_filename, sys.exc_info()[1]))
    
    def pull(self, local_filename, remote_filename):
        
        # determine if the remote filename is actually a directory to place into
        if remote_filename.endswith(support.os_path_separator):
            remote_filename = remote_filename + os.path.basename(local_filename)
        
        # get a list of files (in case it is a dir)
        remote_dirs_to_create = []
        files_to_send = {}
        
        if os.path.isdir(local_filename):
            
            local_dir = os.path.abspath(local_filename)
            remote_dir = remote_filename # TODO: absolute path on remote?  paramiko may allow it via sftp
            
            def directory_visitor(arg, local_currdir, filenames):
                # parse the dir_mapping arg
                local_dir, remote_dir = arg
                
                # figure out the mapped directory
                remote_currdir = os.path.normpath(local_currdir.replace(local_dir, remote_dir))
                remote_dirs_to_create.append(remote_currdir)
                
                for filename in filenames:
                    local_filename = os.path.join(local_currdir, filename)
                    if os.path.isfile(local_filename):
                        remote_filename = os.path.normpath(local_filename.replace(local_dir, remote_dir))
                        files_to_send[local_filename] = remote_filename
            
            # walk the directory structure
            os.path.walk(local_dir, directory_visitor, (local_dir, remote_dir))
            
        else:
            
            files_to_send[local_filename] = remote_filename
        
        # successfully pulled all send specifications, execute them all at once
        # TODO: parallelize this?

        sftp_client = self.ssh_client.open_sftp()
        properly_ordered_remote_dirs_to_create = sorted(remote_dirs_to_create, lambda x,y: len(x) > len(y))

        for new_dir in properly_ordered_remote_dirs_to_create:
            self._mkdir(sftp_client, new_dir)

        for local_filename in files_to_send:
            remote_filename = files_to_send[local_filename]
            self._put(sftp_client, local_filename, remote_filename)


class RemoteSuExpecterStrategy(ExpecterStrategy):

    def __init__(self, user):
        super(RemoteSuExpecterStrategy, self).__init__(user)
        self.ssh_client = paramiko.SSHClient()
        self.ssh_client.load_system_host_keys()
        try:
            self.ssh_client.connect(hostname = user.su_from.host, username = user.su_from.name, password = user.su_from.password)
        except:
            raise AuthenticationFailed("couldn't su for user '%s' via '%s' since login failed (reason: %s)" % (self.user.fullname, self.user.su_from.fullname, sys.exc_info()[1]))
    
    def exec_cmd(self, cmd):
        # modify the command to run via su
        # FIXME: Issue 17: need to escape the command
        if "'" in cmd:
            raise ValueError("RemoteSuExpecterStrategy does not yet support the command '%s'" % cmd)
        cmd = "su %s -c '%s'" % (self.user.name, cmd)
        
        # open an SSH channel (acts like a socket)
        channel = self.ssh_client.get_transport().open_session()
        channel.get_pty()
        
        # grab the streams on that channel
        bufsize = -1
        stdin = channel.makefile('wb', bufsize)
        stdout = channel.makefile('rb', bufsize)
        stderr = channel.makefile_stderr('rb', bufsize)

        # create a callback that will execute the command and return its status
        def exec_and_return_status_code():
            self.log_cmd(cmd)
            channel.exec_command(cmd)
            return channel.recv_exit_status()

        # use the helper to monitor the command
        su_password_prompt = "Password:.*"
        self.respond_with_password(su_password_prompt, self.user.password)
        def kill_this_process():
            channel.close()
        response = self._watch_exec_stream_and_expect_input(cmd, stdin, stdout, stderr, exec_and_return_status_code, kill_this_process)
        self.stop_responding_to("Password:.*")
        return response
        